#include <stdio.h>
#include <stdlib.h>
#include <sqlite3.h>
#include <pthread.h>
#include <time.h>
#include <unistd.h>
#include <string.h>

#define NAME_LEN 128

char db_name_origin[NAME_LEN];
char db_name[NAME_LEN+4];
char tbl_name[NAME_LEN];

// 性能统计结构体
typedef struct {
    double elapsed;      // 耗时(秒)
    double tps;          // 每秒事务数
    long records;        // 操作记录数
    long errors;         // 错误数
} PerfResult;

// 线程参数结构体
typedef struct {
    int thread_id;
    long operations;     // 每个线程的操作数
} ThreadArgs;

// 配置WAL模式与缓存优化
void configure_db(sqlite3 *db, int is_writer) {
    sqlite3_exec(db, "PRAGMA journal_mode=WAL;", NULL, NULL, NULL);
    sqlite3_exec(db, "PRAGMA page_size=4096;", NULL, NULL, NULL);
    sqlite3_exec(db, "PRAGMA cache_size=-20000;", NULL, NULL, NULL); // 20MB缓存
    
    if (is_writer) {
        sqlite3_exec(db, "PRAGMA synchronous=NORMAL;", NULL, NULL, NULL);
    } else {
        sqlite3_exec(db, "PRAGMA synchronous=OFF;", NULL, NULL, NULL);
    }
    
    sqlite3_exec(db, "PRAGMA temp_store=MEMORY;", NULL, NULL, NULL);
    sqlite3_exec(db, "PRAGMA locking_mode=EXCLUSIVE;", NULL, NULL, NULL);
    sqlite3_exec(db, "PRAGMA automatic_index=0;", NULL, NULL, NULL);
}

// 批量写入测试
PerfResult test_bulk_insert(sqlite3 *db, long record_count) {
    sqlite3_stmt *stmt;
    char *data_buf = malloc(32);
    struct timespec start, end;
    int rc;
    char sql[512];
    snprintf(sql, sizeof(sql), "INSERT INTO %s(data) VALUES(?)", tbl_name);
    rc = sqlite3_prepare_v2(db, sql, -1, &stmt, NULL);
    if (rc != SQLITE_OK) {
        fprintf(stderr, "准备插入语句失败: %s\n", sqlite3_errmsg(db));
        free(data_buf);
        return (PerfResult){0, 0, 0, 1};
    }
    
    sqlite3_exec(db, "BEGIN EXCLUSIVE;", NULL, NULL, NULL);
    
    clock_gettime(CLOCK_MONOTONIC, &start);
    
    const long chunk_size = 100000;
    long remaining = record_count;
    
    while (remaining > 0) {
        long batch_size = remaining < chunk_size ? remaining : chunk_size;
        
        for (long i = 0; i < batch_size; i++) {
            snprintf(data_buf, 32, "bulk_data_%ld", i);
            sqlite3_bind_text(stmt, 1, data_buf, -1, SQLITE_STATIC);
            
            if (sqlite3_step(stmt) != SQLITE_DONE) {
                fprintf(stderr, "插入失败: %s\n", sqlite3_errmsg(db));
                break;
            }
            sqlite3_reset(stmt);
        }
        
        sqlite3_exec(db, "COMMIT;", NULL, NULL, NULL);
        if (remaining > chunk_size) {
            sqlite3_exec(db, "BEGIN EXCLUSIVE;", NULL, NULL, NULL);
        }
        
        remaining -= batch_size;
    }
    
    clock_gettime(CLOCK_MONOTONIC, &end);
    
    sqlite3_finalize(stmt);
    free(data_buf);
    
    PerfResult res = {
        .elapsed = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9,
        .records = record_count,
        .errors = 0
    };
    res.tps = res.elapsed > 0 ? record_count / res.elapsed : 0;
    return res;
}

// 查询性能测试
PerfResult test_index_query(sqlite3 *db, long iterations) {
    sqlite3_stmt *stmt;
    struct timespec start, end;
    int rc;
    long found = 0;
    char sql[512];
    snprintf(sql, sizeof(sql), "SELECT data FROM %s WHERE id=?", tbl_name);
    rc = sqlite3_prepare_v2(db, sql, -1, &stmt, NULL);
    if (rc != SQLITE_OK) {
        fprintf(stderr, "准备查询语句失败: %s\n", sqlite3_errmsg(db));
        return (PerfResult){0, 0, 0, 1};
    }
    
    long *ids = malloc(iterations * sizeof(long));
    for (long i = 0; i < iterations; i++) {
        ids[i] = (rand() % 1000000) + 1;
    }
    
    clock_gettime(CLOCK_MONOTONIC, &start);
    
    sqlite3_exec(db, "PRAGMA count_changes=0;", NULL, NULL, NULL);
    
    for (long i = 0; i < iterations; i++) {
        sqlite3_bind_int64(stmt, 1, ids[i]);
        rc = sqlite3_step(stmt);
        if (rc == SQLITE_ROW) {
            found++;
        }
        sqlite3_reset(stmt);
    }
    
    clock_gettime(CLOCK_MONOTONIC, &end);
    
    sqlite3_finalize(stmt);
    free(ids);
    
    PerfResult res = {
        .elapsed = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9,
        .records = iterations,
        .errors = 0
    };
    res.tps = res.elapsed > 0 ? iterations / res.elapsed : 0;
    
    printf("查询命中率: %.2f%%\n", (found * 100.0) / iterations);
    return res;
}

// 线程函数 (并发测试)
void* thread_write(void *arg) {
    ThreadArgs *args = (ThreadArgs*)arg;
    sqlite3 *db;
    char sql[256];
    int rc;
    
    rc = sqlite3_open_v2(db_name, &db, 
                        SQLITE_OPEN_READWRITE | SQLITE_OPEN_FULLMUTEX | SQLITE_OPEN_SHAREDCACHE, 
                        NULL);
    if (rc != SQLITE_OK) {
        fprintf(stderr, "线程 %d 打开数据库失败: %s\n", args->thread_id, sqlite3_errmsg(db));
        return NULL;
    }
    
    configure_db(db, 1);
    
    long *ids = malloc(args->operations * sizeof(long));
    for (long i = 0; i < args->operations; i++) {
        ids[i] = (rand() % 1000000) + 1;
    }
    
    const int batch_size = 100;
    long remaining = args->operations;
    long pos = 0;
    
    while (remaining > 0) {
        int current_batch = remaining < batch_size ? remaining : batch_size;
        sqlite3_exec(db, "BEGIN;", NULL, NULL, NULL);
        
        for (int i = 0; i < current_batch; i++) {
            snprintf(sql, sizeof(sql), "UPDATE %s SET data='updated_%d' WHERE id=%ld;", 
                    tbl_name, args->thread_id, ids[pos++]);
            sqlite3_exec(db, sql, NULL, NULL, NULL);
        }
        
        sqlite3_exec(db, "COMMIT;", NULL, NULL, NULL);
        remaining -= current_batch;
    }
    
    free(ids);
    sqlite3_close(db);
    return NULL;
}

// 并发写入测试
PerfResult test_concurrency(int thread_count, long ops_per_thread) {
    pthread_t threads[thread_count];
    ThreadArgs args[thread_count];
    struct timespec start, end;
    
    for (int i = 0; i < thread_count; i++) {
        args[i].thread_id = i;
        args[i].operations = ops_per_thread;
    }
    
    clock_gettime(CLOCK_MONOTONIC, &start);
    
    for (int i = 0; i < thread_count; i++) {
        if (pthread_create(&threads[i], NULL, thread_write, &args[i]) != 0) {
            fprintf(stderr, "创建线程 %d 失败\n", i);
            return (PerfResult){0, 0, 0, 1};
        }
    }
    
    for (int i = 0; i < thread_count; i++) {
        pthread_join(threads[i], NULL);
    }
    
    clock_gettime(CLOCK_MONOTONIC, &end);
    
    long total_ops = thread_count * ops_per_thread;
    PerfResult res = {
        .elapsed = (end.tv_sec - start.tv_sec) + (end.tv_nsec - start.tv_nsec) / 1e9,
        .records = total_ops,
        .errors = 0
    };
    res.tps = res.elapsed > 0 ? total_ops / res.elapsed : 0;
    return res;
}

// 初始化数据库
int init_database() {
    printf("请输入数据库文件名（如 database）: ");
    while (scanf("%127s", db_name_origin) != 1) {
        fprintf(stderr, "输入错误\n");
    }
    snprintf(db_name, sizeof(db_name), "%s.db", db_name_origin);
    printf("请输入表名（如 users）: ");
    while (scanf("%127s", tbl_name) != 1) {
        fprintf(stderr, "输入错误\n");
    }
    char sql[512];
    snprintf(sql, sizeof(sql),
             "CREATE TABLE IF NOT EXISTS %s ("
             "id INTEGER PRIMARY KEY, "
             "data TEXT NOT NULL);",
             tbl_name);
    sqlite3 *db;
    int rc = sqlite3_open(db_name, &db);
    if (rc != SQLITE_OK) {
        fprintf(stderr, "无法打开数据库: %s\n", sqlite3_errmsg(db));
        sqlite3_close(db);
        return 0;
    }
    
    if (sqlite3_exec(db, sql, NULL, NULL, NULL) != SQLITE_OK) {
        fprintf(stderr, "创建表失败: %s\n", sqlite3_errmsg(db));
        sqlite3_close(db);
        return 0;
    }
    
    configure_db(db, 1);
    sqlite3_close(db);
    return 1;
}

int main() {
    int choice;
    long count;
    srand(time(NULL));
    
    if (!init_database()) {
        return 1;
    }
    
    printf("SQLite功能性能测试工具-V1.0\n");
    printf("1. 批量插入测试\n");
    printf("2. 主键查询测试\n");
    printf("3. 并发写入测试\n");
    printf("请选择测试类型 (1-3): ");
    
    // 验证输入是否为有效整数
    if (scanf("%d", &choice) != 1) {
        fprintf(stderr, "错误：请输入数字1-3\n");
        // 清空输入缓冲区
        while (getchar() != '\n');
        return 1;
    }
    
    sqlite3 *db;
    if (sqlite3_open(db_name, &db) != SQLITE_OK) {
        fprintf(stderr, "无法打开数据库: %s\n", sqlite3_errmsg(db));
        return 1;
    }
    
    switch (choice) { 
        case 1:
            // printf("请输入要插入的记录数: ");
            // scanf("%ld", &count);
            // if (count <= 0) 
            printf("输入数据条数...\n");
            while (scanf("%ld", &count) != 1||count<=0){
              fprintf(stderr, "错误：插入数据条数为大于零的数字，请重新输入：\n");
              // 清空输入缓冲区
              while (getchar() != '\n');
            }
            printf("1. 批量插入测试：数据库名称：%s, 表名称：%s, 插入数据条数：%ld\n",db_name, tbl_name, count);
            printf("是否开始批量插入测试：1 确定，0 取消\n");
            int flag1;
            scanf("%d", &flag1);
            if(flag1==0){
              break;
            }
            printf("开始批量插入测试...\n");
            PerfResult bulk_res = test_bulk_insert(db, count);
            printf("批量插入测试结果: \n");
            printf("插入记录数: %ld\n", bulk_res.records);
            printf("耗时: %.2fs\n", bulk_res.elapsed);
            printf("TPS: %.0f\n", bulk_res.tps);
            break;
            
        case 2:
            // printf("请输入查询次数: ");
            // scanf("%ld", &count);
            // if (count <= 0) 
            printf("输入查询次数...\n");
            while (scanf("%ld", &count) != 1||count<=0){
              fprintf(stderr, "错误：查询次数为大于零的数字，请重新输入：\n");
              // 清空输入缓冲区
              while (getchar() != '\n');
            }
            printf("2. 主键查询测试：数据库名称：%s, 表名称：%s, 查询次数：%ld\n",db_name, tbl_name, count);
            printf("是否开始主键查询测试：1 确定，0 取消\n");
            int flag2;
            scanf("%d", &flag2);
            if(flag2==0){
              break;
            }
            char sql[512];
            snprintf(sql, sizeof(sql), "CREATE INDEX IF NOT EXISTS idx_data ON %s(data)", tbl_name);
            sqlite3_exec(db, sql, NULL, NULL, NULL);
            
            printf("开始主键查询测试...\n");
            PerfResult query_res = test_index_query(db, count);
            printf("主键查询测试结果: \n");
            printf("查询次数: %ld\n", query_res.records);
            printf("耗时: %.4fs\n", query_res.elapsed);
            printf("QPS: %.0f\n", query_res.tps);
            break;
            
        case 3: {
            long ops_per_thread;
            // printf("请输入并发线程数: ");
            // scanf("%ld", &count);
            // if (count <= 0 || count > 100) 
            printf("输入测试线程数...\n");
            while (scanf("%ld", &count) != 1||count<=0){
              fprintf(stderr, "错误：线程数为大于零的数字，请重新输入：\n");
              // 清空输入缓冲区
              while (getchar() != '\n');
            }
            // printf("请输入每个线程的操作数: ");
            // scanf("%ld", &ops_per_thread);
            // if (ops_per_thread <= 0) ops_per_thread = 1000;
            ops_per_thread = 1;
            printf("输入每个线程的更新次数...\n");
            while (scanf("%ld", &ops_per_thread) != 1||ops_per_thread<=0){
              fprintf(stderr, "错误：每个线程的更新次数为大于零的数字，请重新输入：\n");
              // 清空输入缓冲区
              while (getchar() != '\n');
            }
            printf("3. 并发更新测试：数据库名称：%s, 表名称：%s, 测试线程个数：%ld,  每个线程的更新次数：%ld\n",db_name, tbl_name, count,ops_per_thread);
            printf("是否开始并发更新测试：1 确定，0 取消\n");
            int flag3;
            scanf("%d", &flag3);
            if(flag3==0){
              break;
            }
            printf("开始并发更新测试...\n");
            PerfResult conc_res = test_concurrency((int)count, ops_per_thread);
            printf("并发更新测试结果: \n");
            printf("线程数: %ld\n", count);
            printf("总操作数: %ld\n", conc_res.records);
            printf("耗时: %.2fs\n", conc_res.elapsed);
            printf("TPS: %.0f\n", conc_res.tps);
            break;
        }
            
        default:
            printf("无效的选择\n");
            break;
    }
    
    sqlite3_close(db);
    return 0;
}
